import pygame
import heapq

# 初始化 Pygame
pygame.init()

# 设置窗口和网格参数
WIDTH, HEIGHT = 600, 600
GRID_SIZE = 10
CELL_SIZE = WIDTH // GRID_SIZE
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Dynamic Manhattan Path with Direction Control")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
GRAY = (200, 200, 200)

# 曼哈顿距离计算函数
def manhattan_distance(p1, p2):
    return abs(p1[0] - p2[0]) + abs(p1[1] - p2[1])

# 计算路径中的转折次数
def count_turns(path):
    if len(path) < 3:
        return 0
    turns = 0
    for i in range(1, len(path) - 1):
        prev = (path[i-1][0] - path[i][0], path[i-1][1] - path[i][1])
        next = (path[i][0] - path[i+1][0], path[i][1] - path[i+1][1])
        if prev != next:
            turns += 1
    return turns

# 修改后的 A* 算法，尝试满足方向控制
def a_star(start, goal, target_turns=3, start_dir=(0, 1), end_dir=(0, -1)):
    open_set = [(0, 0, start, [start])]  # (f_score, turns, position, path)
    visited = set()
    fallback_path = []  # 备用路径，防止无路径返回

    while open_set:
        f_score, turns, current, path = heapq.heappop(open_set)
        
        # 如果到达终点，检查是否满足条件
        if current == goal:
            if turns == target_turns and len(path) >= 2 and (path[-2][0] - current[0], path[-2][1] - current[1]) == end_dir:
                return path  # 满足所有条件的路径
            elif not fallback_path:  # 保存第一个可行路径作为备用
                fallback_path = path

        if (current, turns) in visited:
            continue
        visited.add((current, turns))
        
        # 方向控制
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]  # 上下左右
        if len(path) == 1:  # 第一步尽量使用 start_dir
            directions = [start_dir] + [d for d in directions if d != start_dir]  # 优先 start_dir，但允许其他方向
        
        for dx, dy in directions:
            neighbor = (current[0] + dx, current[1] + dy)
            
            if (0 <= neighbor[0] < GRID_SIZE and 
                0 <= neighbor[1] < GRID_SIZE and 
                neighbor not in path):
                
                new_path = path + [neighbor]
                new_turns = count_turns(new_path)
                
                # 如果转折次数超过目标，跳过
                if new_turns > target_turns:
                    continue
                
                if (neighbor, new_turns) not in visited:
                    g_score = len(new_path) - 1
                    h_score = manhattan_distance(neighbor, goal)
                    f_score = g_score + h_score
                    heapq.heappush(open_set, (f_score, new_turns, neighbor, new_path))
    
    return fallback_path if fallback_path else []  # 如果没有完美路径，返回备用路径

# 绘制网格和路径（使用线段）
def draw_grid(start, goal, path):
    screen.fill(WHITE)
    
    # 绘制网格线
    for x in range(0, WIDTH, CELL_SIZE):
        pygame.draw.line(screen, GRAY, (x, 0), (x, HEIGHT))
    for y in range(0, HEIGHT, CELL_SIZE):
        pygame.draw.line(screen, GRAY, (0, y), (WIDTH, y))
    
    # 绘制路径线段
    if path:
        for i in range(len(path) - 1):
            start_pos = (path[i][1] * CELL_SIZE + CELL_SIZE // 2, path[i][0] * CELL_SIZE + CELL_SIZE // 2)
            end_pos = (path[i+1][1] * CELL_SIZE + CELL_SIZE // 2, path[i+1][0] * CELL_SIZE + CELL_SIZE // 2)
            pygame.draw.line(screen, RED, start_pos, end_pos, 3)  # 线宽为 3
    
    # 绘制起点和终点
    pygame.draw.circle(screen, BLACK, (start[1] * CELL_SIZE + CELL_SIZE // 2, start[0] * CELL_SIZE + CELL_SIZE // 2), CELL_SIZE // 3)
    pygame.draw.circle(screen, GREEN, (goal[1] * CELL_SIZE + CELL_SIZE // 2, goal[0] * CELL_SIZE + CELL_SIZE // 2), CELL_SIZE // 3)

# 主程序
def main():
    # 固定起点
    start = (GRID_SIZE // 2, GRID_SIZE // 2)  # 起点设为网格中心
    running = True
    target_turns = 3  # 3 次转折 = 4 段折线
    start_dir = (0, 1)  # 起点向右离开
    end_dir = (0, -1)  # 终点从下进入
    
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
        
        # 获取鼠标位置并转换为网格坐标
        mouse_pos = pygame.mouse.get_pos()
        goal = (mouse_pos[1] // CELL_SIZE, mouse_pos[0] // CELL_SIZE)
        
        # 确保终点在网格范围内
        goal = (max(0, min(goal[0], GRID_SIZE-1)), max(0, min(goal[1], GRID_SIZE-1)))
        
        # 计算路径
        path = a_star(start, goal, target_turns, start_dir, end_dir)
        
        # 绘制网格和路径
        draw_grid(start, goal, path)
        pygame.display.flip()
    
    pygame.quit()

if __name__ == "__main__":
    main()